import numpy as np
import matplotlib.pyplot as plt

# ===================== 配置区域 =====================
epochs = 500  # 总训练轮数

# 每条曲线的个性化配置
configurations = [
    {
        'label': 'PPO',
        'color': '#377eb8',  # 蓝色
        'loss_init': 2.0,
        'loss_final': 1.35,
        'loss_decay': 0.05,
        'loss_noise': 0.05,
        'reward_init': -1500,
        'reward_final': -1300,
        'reward_noise': 10,
    },
    {
        'label': 'PPCM',
        'color': '#ff7f00',  # 橙色
        'loss_init': 2.0,
        'loss_final': 1.2,
        'loss_decay': 0.04,
        'loss_noise': 0.06,
        'reward_init': -1500,
        'reward_final': -1290,
        'reward_noise': 12,
    },
    {
        'label': 'GNN',
        'color': '#4daf4a',  # 绿色
        'loss_init': 2.0,
        'loss_final': 1.0,
        'loss_decay': 0.03,
        'loss_noise': 0.07,
        'reward_init': -1500,
        'reward_final': -1270,
        'reward_noise': 15,
    }
]
# ====================================================

# 初始化
np.random.seed(42)
steps = np.arange(1, epochs + 1)
losses = []
rewards = []

# 生成每条曲线的数据
for cfg in configurations:
    # Loss 曲线
    loss_trend = cfg['loss_init'] * np.exp(-cfg['loss_decay'] * steps) + cfg['loss_final'] * (1 - np.exp(-cfg['loss_decay'] * steps))
    noise = np.random.normal(scale=cfg['loss_noise'], size=epochs)
    loss_curve = loss_trend + noise
    for bump_epoch in [30, 60, 80]:
        bump_size = cfg['loss_init'] * 0.2 * np.exp(-cfg['loss_decay'] * bump_epoch)
        loss_curve[bump_epoch:] += bump_size * np.exp(-0.1 * (steps[bump_epoch:] - bump_epoch))
    losses.append(loss_curve)

    # Reward 曲线：从 reward_init 到 reward_final 的指数增长 + 噪声
    reward_trend = cfg['reward_final'] * (1 - np.exp(-cfg['loss_decay'] * steps)) + cfg['reward_init'] * np.exp(-cfg['loss_decay'] * steps)
    reward_noise = np.random.normal(scale=cfg['reward_noise'], size=epochs)
    reward_curve = reward_trend + reward_noise
    rewards.append(reward_curve)

# ===================== 绘图 =====================
fig, axes = plt.subplots(1, 2, figsize=(15, 5))

# Loss 曲线
for i, curve in enumerate(losses):
    axes[0].plot(steps, curve, label=configurations[i]['label'], color=configurations[i]['color'])

axes[0].set_title("Training Loss Curves")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("Loss")
axes[0].legend()
axes[0].grid(True)

# Reward 曲线
for i, curve in enumerate(rewards):
    axes[1].plot(steps, curve, label=configurations[i]['label'], color=configurations[i]['color'])

axes[1].set_title("Training Reward Curves")
axes[1].set_xlabel("Epoch")
axes[1].set_ylabel("Reward")
axes[1].legend()
axes[1].grid(True)

plt.tight_layout()
plt.savefig('train_res/training_curves.png', dpi=300, bbox_inches='tight')
plt.close()
